{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型加载与保存"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "训练前:  tensor([[[ 0.0663,  0.1155,  0.0588, -0.0793, -0.0634],\n",
      "         [-0.0235,  0.1004,  0.0122, -0.0956,  0.0776],\n",
      "         [ 0.0357, -0.1122,  0.0310,  0.0453, -0.0823],\n",
      "         [-0.1140, -0.0801, -0.1007, -0.0095, -0.0445],\n",
      "         [-0.0524, -0.0345,  0.0624,  0.0884, -0.0357]],\n",
      "\n",
      "        [[-0.0626,  0.0563,  0.0243, -0.1103,  0.0023],\n",
      "         [ 0.0951, -0.0181, -0.1153, -0.0853,  0.0355],\n",
      "         [ 0.0713, -0.0709, -0.0379,  0.0242,  0.0476],\n",
      "         [-0.0903,  0.0322,  0.0994,  0.0152, -0.0590],\n",
      "         [ 0.0324,  0.0886,  0.0415,  0.0373, -0.0832]],\n",
      "\n",
      "        [[-0.0949,  0.0828, -0.0199, -0.0991, -0.0592],\n",
      "         [ 0.0776, -0.0896,  0.0081, -0.0105, -0.0269],\n",
      "         [ 0.1046, -0.0198, -0.0831,  0.0023,  0.0971],\n",
      "         [-0.0556, -0.0557, -0.1149,  0.0150,  0.0603],\n",
      "         [-0.0686,  0.0720, -0.0763, -0.0072,  0.1149]]],\n",
      "       grad_fn=<SelectBackward0>)\n",
      "训练后:  tensor([[[2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.]],\n",
      "\n",
      "        [[2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.]],\n",
      "\n",
      "        [[2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.]]], grad_fn=<SelectBackward0>)\n"
     ]
    }
   ],
   "source": [
    "class LeNet2(nn.Module):\n",
    "    def __init__(self, classes):\n",
    "        super(LeNet2, self).__init__()\n",
    "        self.features = nn.Sequential(\n",
    "            nn.Conv2d(3, 6, 5),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Conv2d(6, 16, 5),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2, 2)\n",
    "        )\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Linear(16*5*5, 120),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(120, 84),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(84, classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.features(x)\n",
    "        x = x.view(x.size()[0], -1)\n",
    "        x = self.classifier(x)\n",
    "        return x\n",
    "\n",
    "    def initialize(self):\n",
    "        for p in self.parameters():\n",
    "            p.data.fill_(2020)\n",
    "\n",
    "net = LeNet2(classes=2019)\n",
    "\n",
    "# \"训练\"\n",
    "print(\"训练前: \", net.features[0].weight[0, ...])\n",
    "net.initialize()\n",
    "print(\"训练后: \", net.features[0].weight[0, ...])\n",
    "\n",
    "path_model = \"./model.pkl\"\n",
    "path_state_dict = \"./model_state_dict.pkl\"\n",
    "\n",
    "# 保存整个模型\n",
    "torch.save(net, path_model)\n",
    "\n",
    "# 保存模型参数\n",
    "net_state_dict = net.state_dict()\n",
    "torch.save(net_state_dict, path_state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LeNet2(\n",
      "  (features): Sequential(\n",
      "    (0): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))\n",
      "    (1): ReLU()\n",
      "    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "    (3): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))\n",
      "    (4): ReLU()\n",
      "    (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "  )\n",
      "  (classifier): Sequential(\n",
      "    (0): Linear(in_features=400, out_features=120, bias=True)\n",
      "    (1): ReLU()\n",
      "    (2): Linear(in_features=120, out_features=84, bias=True)\n",
      "    (3): ReLU()\n",
      "    (4): Linear(in_features=84, out_features=2019, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "path_model = \"./model.pkl\"\n",
    "net_load = torch.load(path_model)\n",
    "\n",
    "print(net_load)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "加载前:  tensor([[[ 0.0775,  0.0374,  0.0163,  0.0196, -0.0884],\n",
      "         [ 0.0293, -0.1051, -0.0362,  0.1122, -0.0616],\n",
      "         [ 0.0083,  0.0274,  0.0158,  0.0301,  0.0937],\n",
      "         [-0.0459, -0.1062,  0.0510, -0.0058,  0.1046],\n",
      "         [-0.0672, -0.0204,  0.0134,  0.0594,  0.0421]],\n",
      "\n",
      "        [[ 0.0058, -0.0435, -0.0550,  0.0591, -0.1067],\n",
      "         [ 0.0929,  0.0202, -0.0027,  0.0264,  0.0409],\n",
      "         [ 0.0038, -0.0219, -0.0522, -0.0065,  0.0717],\n",
      "         [-0.0300, -0.0819, -0.0238, -0.0132, -0.0364],\n",
      "         [ 0.0258, -0.0238, -0.0680, -0.0172,  0.0902]],\n",
      "\n",
      "        [[-0.1087,  0.0948, -0.0848,  0.1148, -0.0212],\n",
      "         [-0.0634,  0.0479,  0.0064, -0.0287,  0.0732],\n",
      "         [-0.1080,  0.0522, -0.0891, -0.1137,  0.0838],\n",
      "         [ 0.0740,  0.0965,  0.0893, -0.1075,  0.0277],\n",
      "         [-0.0060, -0.0713,  0.0996,  0.0865, -0.0181]]],\n",
      "       grad_fn=<SelectBackward0>)\n",
      "加载后:  tensor([[[2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.]],\n",
      "\n",
      "        [[2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.]],\n",
      "\n",
      "        [[2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.],\n",
      "         [2020., 2020., 2020., 2020., 2020.]]], grad_fn=<SelectBackward0>)\n"
     ]
    }
   ],
   "source": [
    "path_state_dict = \"./model_state_dict.pkl\"\n",
    "state_dict_load = torch.load(path_state_dict)\n",
    "net_new = LeNet2(classes=2019)\n",
    "\n",
    "print(\"加载前: \", net_new.features[0].weight[0, ...])\n",
    "net_new.load_state_dict(state_dict_load)\n",
    "print(\"加载后: \", net_new.features[0].weight[0, ...])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model.state_dict:\n",
      "conv1.weight \t torch.Size([6, 3, 5, 5])\n",
      "conv1.bias \t torch.Size([6])\n",
      "conv2.weight \t torch.Size([16, 6, 5, 5])\n",
      "conv2.bias \t torch.Size([16])\n",
      "fc1.weight \t torch.Size([120, 400])\n",
      "fc1.bias \t torch.Size([120])\n",
      "fc2.weight \t torch.Size([84, 120])\n",
      "fc2.bias \t torch.Size([84])\n",
      "fc3.weight \t torch.Size([10, 84])\n",
      "fc3.bias \t torch.Size([10])\n",
      "Optimizer`s state_dict:\n",
      "state \t {}\n",
      "param_groups \t [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'maximize': False, 'foreach': None, 'differentiable': False, 'params': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]}]\n"
     ]
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    " \n",
    "#define model\n",
    "class TheModelClass(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(TheModelClass,self).__init__()\n",
    "        self.conv1=nn.Conv2d(3,6,5)\n",
    "        self.pool=nn.MaxPool2d(2,2)\n",
    "        self.conv2=nn.Conv2d(6,16,5)\n",
    "        self.fc1=nn.Linear(16*5*5,120)\n",
    "        self.fc2=nn.Linear(120,84)\n",
    "        self.fc3=nn.Linear(84,10)\n",
    " \n",
    "    def forward(self,x):\n",
    "        x=self.pool(F.relu(self.conv1(x)))\n",
    "        x=self.pool(F.relu(self.conv2(x)))\n",
    "        x=x.view(-1,16*5*5)\n",
    "        x=F.relu(self.fc1(x))\n",
    "        x=F.relu(self.fc2(x))\n",
    "        x=self.fc3(x)\n",
    "        return x\n",
    " \n",
    "def main():\n",
    "    # Initialize model\n",
    "    model = TheModelClass()\n",
    " \n",
    "    #Initialize optimizer\n",
    "    optimizer=optim.SGD(model.parameters(),lr=0.001,momentum=0.9)\n",
    " \n",
    "    #print model's state_dict\n",
    "    print('Model.state_dict:')\n",
    "    for param_tensor in model.state_dict():\n",
    "        #打印 key value字典\n",
    "        print(param_tensor,'\\t',model.state_dict()[param_tensor].size())\n",
    " \n",
    "    #print optimizer's state_dict\n",
    "    print('Optimizer`s state_dict:')\n",
    "    for var_name in optimizer.state_dict():\n",
    "        print(var_name,'\\t',optimizer.state_dict()[var_name])\n",
    " \n",
    " \n",
    " \n",
    "if __name__=='__main__':\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
